import re
import matplotlib.pyplot as plt
import argparse  # 导入用于处理命令行参数的库
import sys


def parse_log_data(log_content):
    """
    从日志内容中解析最终的汇总数据。
    支持两种格式，并分别返回它们的结果。
    """

    # 格式 1: "INFO - MODELNAME | AUC: ... | Acc: ..."
    pattern1 = (
        r"INFO - ([\w]+)\s+\| AUC: ([\d\.]+) ± ([\d\.]+) \| Acc: ([\d\.]+) ± ([\d\.]+)"
    )

    # 格式 2: "INFO - 模型: MODELNAME \n ... 平均 最佳Val-AUC: ... \n ... 平均 最佳Val-Acc: ..."
    pattern2 = r"INFO - 模型: ([\w]+).*?平均 最佳Val-AUC: ([\d\.]+) ± ([\d\.]+).*?平均 最佳Val-Acc: ([\d\.]+) ± ([\d\.]+)"

    matches1 = re.findall(pattern1, log_content)
    matches2 = re.findall(pattern2, log_content, re.DOTALL)

    def process_matches(matches):
        """辅助函数，用于将匹配列表转换为数据列表"""
        models = []
        auc_means = []
        auc_stds = []
        acc_means = []
        acc_stds = []
        for match in matches:
            models.append(match[0])
            auc_means.append(float(match[1]))
            auc_stds.append(float(match[2]))
            acc_means.append(float(match[3]))
            acc_stds.append(float(match[4]))
        return models, auc_means, auc_stds, acc_means, acc_stds

    # 分别处理并返回两组数据
    data_format1 = process_matches(matches1)
    data_format2 = process_matches(matches2)

    return data_format1, data_format2


def plot_results(
    models, auc_means, auc_stds, acc_means, acc_stds, plot_title, output_filename
):
    """
    使用麦黄色系绘制结果图表 (白色背景)。
    所有文本均为英文。
    """

    # --- 麦黄色系 ---
    # 背景色 (白色)
    bg_color = "#FFFFFF"

    # 条形图颜色 (小麦色系)
    wheat_palette = [
        "#F5DEB3",  # Wheat (小麦色)
        "#DEB887",  # BurlyWood (硬木色)
        "#FFDAB9",  # PeachPuff (桃色)
        "#D2B48C",  # Tan (棕褐色)
        "#F4A460",  # SandyBrown (沙棕色)
        "#FFE4B5",  # Moccasin (鹿皮色)
    ]

    # 边框和网格线 (暗金色)
    edge_color = "#DAA520"

    # 文字和误差线 (黑色)
    text_color = "black"
    # ---

    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 8))
    fig.patch.set_facecolor(bg_color)

    x_pos = range(len(models))

    # 为每个bar生成颜色列表
    colors_list = [wheat_palette[i % len(wheat_palette)] for i in x_pos]

    # --- 绘制 AUC 图表 ---
    ax1.bar(
        x_pos,
        auc_means,
        yerr=auc_stds,
        align="center",
        alpha=0.9,
        color=colors_list,
        edgecolor=edge_color,
        linewidth=1.5,
        capsize=7,
        ecolor=text_color,
    )

    ax1.set_facecolor(bg_color)
    ax1.set_ylabel("AUC (Area Under the Curve)", color=text_color, fontsize=12)
    ax1.set_xticks(x_pos)
    ax1.set_xticklabels(models, color=text_color, fontsize=11)
    ax1.set_title("Model AUC Mean", color=text_color, fontsize=16, pad=20)

    min_auc = min(m - s for m, s in zip(auc_means, auc_stds))
    max_auc = max(m + s for m, s in zip(auc_means, auc_stds))
    ax1.set_ylim([max(0, min_auc - 0.1), min(1, max_auc + 0.1)])

    # 在条形图上添加数值标签 (AUC)
    for i, (val, err) in enumerate(zip(auc_means, auc_stds)):
        ax1.text(
            i,
            val + err + 0.01,
            f"{val:.4f}",
            ha="center",
            va="bottom",
            color=text_color,
            fontsize=10,
        )

    # --- 绘制 Accuracy 图表 ---
    ax2.bar(
        x_pos,
        acc_means,
        yerr=acc_stds,
        align="center",
        alpha=0.9,
        color=colors_list,
        edgecolor=edge_color,
        linewidth=1.5,
        capsize=7,
        ecolor=text_color,
    )

    ax2.set_facecolor(bg_color)
    ax2.set_ylabel("Acc (Accuracy)", color=text_color, fontsize=12)
    ax2.set_xticks(x_pos)
    ax2.set_xticklabels(models, color=text_color, fontsize=11)
    ax2.set_title("Model Accuracy Mean", color=text_color, fontsize=16, pad=20)

    min_acc = min(m - s for m, s in zip(acc_means, acc_stds))
    max_acc = max(m + s for m, s in zip(acc_means, acc_stds))
    ax2.set_ylim([max(0, min_acc - 0.1), min(1, max_acc + 0.1)])

    # 在条形图上添加数值标签 (Acc)
    for i, (val, err) in enumerate(zip(acc_means, acc_stds)):
        ax2.text(
            i,
            val + err + 0.01,
            f"{val:.4f}",
            ha="center",
            va="bottom",
            color=text_color,
            fontsize=10,
        )

    # --- 通用样式调整 ---
    for ax in [ax1, ax2]:
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.spines["bottom"].set_color(text_color)
        ax.spines["left"].set_color(text_color)
        ax.tick_params(axis="x", colors=text_color)
        ax.tick_params(axis="y", colors=text_color)
        ax.yaxis.grid(True, linestyle="--", which="major", color=edge_color, alpha=0.5)
        ax.set_axisbelow(True)

    fig.suptitle(plot_title, fontsize=20, color=text_color, y=0.98)
    plt.tight_layout(rect=[0, 0.03, 1, 0.95])

    plt.savefig(output_filename, dpi=300, facecolor=bg_color)
    print(f"Chart saved successfully to: {output_filename}")


def read_log_file(filepath):
    """
    从指定路径读取日志文件内容。
    """
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            return f.read()
    except FileNotFoundError:
        print(f"Error: File not found {filepath}", file=sys.stderr)
        return None
    except Exception as e:
        print(f"Error reading file: {e}", file=sys.stderr)
        return None


def main():

    parser = argparse.ArgumentParser(
        description="Parse model performance data from log files and plot charts."
    )
    parser.add_argument(
        "logfile",
        nargs="?",
        default="/data0/lcy/trident/fpt/logs/20251030_0000_cross_mil_val_training.log",
        help="Path to the log file to parse.",
    )
    args = parser.parse_args()
    logfile = args.logfile

    # 2. 读取日志文件
    print(f"Reading log from {logfile}...")
    log_data = read_log_file(logfile)

    if log_data is None:
        return  # 如果文件读取失败，则退出

    # 3. 解析数据 (拆分为两种格式)
    (data1), (data2) = parse_log_data(log_data)

    models1, auc_means1, auc_stds1, acc_means1, acc_stds1 = data1
    models2, auc_means2, auc_stds2, acc_means2, acc_stds2 = data2

    # 4. 为格式1绘图 (如果有数据)
    if models1:
        print(f"Successfully parsed {len(models1)} models for format 1: {models1}")
        plot_results(
            models1,
            auc_means1,
            auc_stds1,
            acc_means1,
            acc_stds1,
            "Log Format 1 Results",
            "log_format_1_results.png",
        )
    else:
        print("No data found matching Log Format 1.")

    # 5. 为格式2绘图 (如果有数据)
    if models2:
        print(f"Successfully parsed {len(models2)} models for format 2: {models2}")
        plot_results(
            models2,
            auc_means2,
            auc_stds2,
            acc_means2,
            acc_stds2,
            "Log Format 2 Results",
            "log_format_2_results.png",
        )
    else:
        print("No data found matching Log Format 2.")

    if not models1 and not models2:
        print(
            "Error: No matching summary data found in log for either format.",
            file=sys.stderr,
        )
        return


if __name__ == "__main__":
    main()
